import ncnn
import cv2
import numpy as np


# 创建ncnn的网络对象
net = ncnn.Net()

# 加载ONNX模型
net.load_param('C:\\Users\\23659\Downloads\\best-sim-opt.param')
net.load_model('C:\\Users\\23659\Downloads\\best-sim-opt.bin')

# 加载图像
image = cv2.imread(r'C:\Users\23659\Desktop\yolov5-master\ss\train\images\1709446496066.jpeg')

# 调整图像尺寸为模型输入尺寸
input_size = (800, 800)
resized_image = cv2.resize(image, input_size)

# 减去均值
mean_vals = (0.37802792*255.0,0.32611448*255.0,0.29480308*255.0)
norm_vals = (1 / 0.348492 / 255.0, 1 / 0.3070657 / 255.0, 1 / 0.28770673 / 255.0)
input_blob = ncnn.Mat.from_pixels(
    resized_image, ncnn.Mat.PixelType.PIXEL_BGR2RGB, 800, 800)
# 运行网络
input_blob.substract_mean_normalize(mean_vals, norm_vals)
ex = net.create_extractor()
# net_input = ncnn.Extractor(net)
ex.input("input", input_blob)
output_blob = ncnn.Mat()
ex.extract("output", output_blob)

# 获取分类结果
# output_data = output_blob.to_numpy()


output_blob = output_blob.reshape(2,800 , 800)
output_blob = np.array(output_blob)
mask = output_blob[0]>0.8
print(800*800,';;;;;',np.sum(mask))

img0 = np.array(image*mask[:,:,None],dtype=np.uint8)

cv2.imshow('hh',img0)
cv2.waitKey(0)


img1 =  np.array(image*~mask[:,:,None],dtype=np.uint8)

cv2.imshow('hh1',img1)
cv2.waitKey(0)

print(1)

